# -*- coding: utf-8 -*-
"""
Created on Sat Dec 12 09:56:51 2020

@author: Administrator
"""

import numpy as np




def bjk(B, j, k):
    return B[j][k]


def fw(A, B, Pi, O, T, N):
    """
    forward algorithm
    """
    # size of a is len of N * T
    alpha = np.zeros([N, T])

    for t in range(T):
        for j in range(N):
            # if t = 0, then initial the value of a
            if t == 0:
                alpha[j][t] = Pi[j] * bjk(B, j, O[t])
            else:
                for i in range(N):
                    #a_t+1(j) = sum of a_t(i)* aij * B_j(O_t+1) 
                    alpha[j][t] += alpha[i][t - 1] * A[i][j] * bjk(B, j, O[t])

    return alpha


def bw(A, B, Pi, O, T, N):
    """
    backward algorithm
    input:
        A, B, Pi: lambda
        O: oberservations
        T: times
        N: number of status
    """
    beta = np.zeros([N, T])
    for t in reversed(range(T)):
        for i in range(N):
            #set initial value = 1
            if t == T-1:
                beta[i][t] = 1
            else:
                for j in range(N):
                    #update b_t(i) by sum of aij * B_j(o_t+1) * b_t+1(j) 
                    beta[i][t] += bjk(B, j, O[t+1]) * beta[j][t+1] * A[i][j]

    return beta


def gamma_t_i(t, alpha, beta, i):
    """
    given O, lambda, Probability of being in state i at time t
    gamma_t(i) = P(it=qi|O, lambda) = P(it=qi, O|lambda) / P(O|lambda)
               = alpha_ti * beta_ti / [aloha_tj + beta_j for j in N].sum()  
    """
    sum_of_t = 0
    for j in range(len(alpha)):
        sum_of_t += alpha[j][t] * beta[j][t]
        
    return alpha[i][t] * beta[i][t] / sum_of_t

def ksi(t, A, B, alpha, beta, i, j):
    """
    given O, lambda, probability of being qi in time t, qj in time t+1
    kessi_tij = P(it=qi, it+1=qj|O, lambda)
              = P(it=qi, it+1=qj, O|lambda) / p(O|lamnda) 
              = Pit=qi, it+1=qj, O|lambda) / ΣΣ(Pit=qi, it+1=qj, O|lambda)
    ati * aij * bj(Ot+1) * beta[j][t+1]
    """
    ksi_i_j = alpha[i][t] * A[i][j] * bjk(B, j, O[t+1])* beta[j][t+1]
    sum_of_t = 0
    for i in range(len(alpha)):
        for j in range(len(alpha)):
            sum_of_t += alpha[i][t] * A[i][j] * bjk(B, j, O[t+1])* beta[j][t+1]

    return ksi_i_j / sum_of_t



T = 3
O = [0, 1, 0]
A = np.ones((3, T))
A = A / A.sum(axis = 1)
B = np.ones((3, 2))
B = B / len(B[0])
N = len(A)
Pi = np.transpose(np.array([0.3, 0.3, 0.4]))

step = 0
while step < 500:
    step += 1
    alpha = fw(A, B, Pi, O, T, N)
    beta = bw(A, B, Pi, O, T, N)
    
    for i in range(N):
        Pi[i] = gamma_t_i(0, alpha, beta, i)
    
    for i in range(N):
        for j in range(N):
            _ksi = 0
            _gamma = 0
            for t in range(T-1):
                _ksi += ksi(t, A, B, alpha, beta, i, j)
                _gamma += gamma_t_i(t, alpha, beta, i)
                
            A[i][j] = _ksi / _gamma
    
    for j in range(N):
        for k in range(2):
            _gamma_numerator = 0
            _gamma_denominator = 0
            for t in range(T):
                if O[t] == k:
                    _gamma_numerator += gamma_t_i(t, alpha, beta, j)
                _gamma_denominator += gamma_t_i(t, alpha, beta, j)
                
            B[j][k] = _gamma_numerator / _gamma_denominator
            